# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import time
import logging
from common.core.cron import CronBase
from common.enums.apply import DocsEnum
from common.enums.apply import MusicEnum
from common.enums.apply import VideoEnum
from common.enums.apply import DrawsEnum
from common.enums.agent import AgentEnum
from common.enums.wallet import WalletEnum
from common.cache.queue_cache import QueueCache
from common.models.users import UserWalletModel
from common.models.docs import DocsFilesModel
from common.models.music import MusicRecordsModel
from common.models.video import VideoRecordsModel
from common.models.draws import DrawsRecordsModel
from common.models.agents import AgentKnowArchiveModel

logger = logging.getLogger(__name__)


class Command(CronBase):
    @classmethod
    async def run(cls, **kwargs):
        cron_index: int = int(kwargs["w_ix"])
        if cron_index != 1:
            return

        try:
            await cls._distribute_em()
            await cls._distribute_qa()
            await cls._distribute_docs()
            await cls._distribute_draws()
            await cls._distribute_video()
            await cls._distribute_music()
        except Exception as e:
            logger.error(f"crontab error: {str(e)}")

    @classmethod
    async def _distribute_em(cls):
        """ 向量训练调度器 """
        try:
            # 队列中元素数量
            queue_data = await QueueCache.queue_peek(QueueCache.EM_JOB)
            if len(queue_data) > 50:
                return None

            # 分片等待中数据
            ck_records = await (AgentKnowArchiveModel
                                .filter(status=AgentEnum.EMB_WAIT)
                                .filter(train_mode="chunk")
                                .filter(is_delete=0)
                                .order_by("id")
                                .limit(200)
                                .values("id"))

            # 拆分等待中数据
            qa_records = await (AgentKnowArchiveModel
                                .filter(train_mode="qa")
                                .filter(status=AgentEnum.EMB_WAIT)
                                .filter(qa_status=AgentEnum.EMB_YES)
                                .filter(is_delete=0)
                                .order_by("id")
                                .limit(200)
                                .values("id"))

            # 推入训练队列
            ids = []
            wait_records = ck_records + qa_records
            for item in wait_records:
                ids.append(item["id"])
                await QueueCache.queue_push(QueueCache.EM_JOB, str(item["id"]))

            # 调整为排队中
            if ids:
                await AgentKnowArchiveModel.filter(id__in=ids).update(
                    status=AgentEnum.EMB_LINE,
                    update_time=int(time.time())
                )

            return None
        except Exception as e:
            logger.error(f"crontab _distribute_em: " + str(e))

    @classmethod
    async def _distribute_qa(cls):
        """ QA拆分调度器 """
        try:
            # 队列中元素数量
            queue_data = await QueueCache.queue_peek(QueueCache.QA_JOB)
            if len(queue_data) > 10:
                return None

            # 获取等待中数据
            wait_records = await (AgentKnowArchiveModel
                                  .filter(qa_status=AgentEnum.EMB_WAIT)
                                  .filter(train_mode="qa")
                                  .filter(is_delete=0)
                                  .order_by("create_time")
                                  .limit(50)
                                  .values("id"))
            # 推入拆分队列
            ids = []
            for item in wait_records:
                ids.append(item["id"])
                await QueueCache.queue_push(QueueCache.QA_JOB, str(item["id"]))

            # 调整为排队中
            if len(ids) > 0:
                await AgentKnowArchiveModel.filter(id__in=ids).update(
                    qa_status=AgentEnum.EMB_LINE,
                    update_time=int(time.time())
                )

            return None
        except Exception as e:
            logger.error(f"crontab _distribute_qa: " + str(e))

    @classmethod
    async def _distribute_docs(cls):
        # 队列中元素数量
        queue_data = await QueueCache.queue_peek(QueueCache.DOCS_JOB)
        if len(queue_data) > 10:
            return None

        # 获取等待中数据
        wait_records = await (DocsFilesModel
                              .filter(status=0)
                              .filter(is_delete=0)
                              .order_by("create_time")
                              .limit(50)
                              .values("id"))

        # 推入训练队列
        ids = []
        for item in wait_records:
            ids.append(int(item["id"]))
            await QueueCache.queue_push(QueueCache.DOCS_JOB, str(item["id"]))

        # 调整为排队中
        if ids:
            await DocsFilesModel.filter(id__in=ids).update(
                status=AgentEnum.EMB_LINE,
                update_time=int(time.time())
            )

        return None

    @classmethod
    async def _distribute_draws(cls):
        try:
            # 队列中元素数量
            queue_data = await QueueCache.queue_peek(QueueCache.DRAW_JOB)
            if len(queue_data) > 10:
                return None

            # 查询超时未完成 (10分钟未处理完)
            expire_time: int = int(10 * 60)
            current_time: int = int(time.time())
            err_records = await (DrawsRecordsModel
                                 .filter(is_delete=0)
                                 .filter(status=DrawsEnum.STATUS_ING)
                                 .filter(start_time__lte=(current_time - expire_time))
                                 .filter(engine__in=["mj", "doubao", "dalle"])
                                 .limit(100)
                                 .order_by("id")
                                 .values("id", "user_id", "use_points", "create_time"))

            # 处理超时未完成
            for record in err_records:
                # 已超时
                await DrawsRecordsModel.filter(id=record["id"]).update(
                    fail_reason="任务响应超时",
                    status=DrawsEnum.STATUS_FAIL,
                    update_time=int(time.time())
                )
                # 回退积分
                if record["use_points"] > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record["user_id"],
                        change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                        change_amount=record["use_points"],
                        project="AI绘画"
                    )

            # 获取等待中数据
            wait_records = await (DrawsRecordsModel
                                  .filter(is_delete=0)
                                  .filter(engine__in=["doubao", "dalle"])
                                  .filter(status=DrawsEnum.STATUS_WAIT)
                                  .limit(50)
                                  .order_by("id")
                                  .values("id"))

            # 处理待绘图任务
            wait_lists = [item["id"] for item in wait_records]
            if wait_lists:
                # 更新数据排队中
                await (DrawsRecordsModel
                       .filter(is_delete=0)
                       .filter(id__in=wait_lists)
                       .update(
                            status=DrawsEnum.STATUS_ING,
                            start_time=int(time.time()),
                            update_time=int(time.time())
                       ))
                # 将数据推入队列
                for val in wait_lists:
                    await QueueCache.queue_push(QueueCache.DRAW_JOB, val)
            return None
        except Exception as e:
            logger.error(f"crontab _distribute_draws: " + str(e))

    @classmethod
    async def _distribute_video(cls):
        try:
            # 队列中元素数量
            queue_data = await QueueCache.queue_peek(QueueCache.DRAW_JOB)
            if len(queue_data) > 10:
                return None

            # 查询超时未完成 (10分钟未处理完)
            expire_time: int = int(10 * 60)
            current_time: int = int(time.time())
            err_records = await (VideoRecordsModel
                                 .filter(is_delete=0)
                                 .filter(status=VideoEnum.STATUS_ING)
                                 .filter(start_time__lte=(current_time - expire_time))
                                 .limit(100)
                                 .order_by("id")
                                 .values("id", "user_id", "use_points", "create_time"))

            # 处理超时未完成
            for record in err_records:
                # 已超时
                await VideoRecordsModel.filter(id=record["id"]).update(
                    fail_reason="任务响应超时",
                    status=VideoEnum.STATUS_FAIL,
                    update_time=int(time.time())
                )
                # 回退积分
                if record["use_points"] > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record["user_id"],
                        change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                        change_amount=record["use_points"],
                        project="AI视频"
                    )

            # 获取等待中数据
            wait_records = await (VideoRecordsModel
                                  .filter(is_delete=0)
                                  .filter(status__in=[VideoEnum.STATUS_WAIT, VideoEnum.STATUS_ING])
                                  .limit(50)
                                  .order_by("id")
                                  .values("id"))

            # 处理待视频任务
            wait_lists = [item["id"] for item in wait_records]
            if wait_lists:
                # 更新数据排队中
                await (VideoRecordsModel
                       .filter(id__in=wait_lists)
                       .filter(is_delete=0)
                       .update(
                            status=VideoEnum.STATUS_ING,
                            start_time=int(time.time()),
                            update_time=int(time.time())
                       ))
                # 将数据推入队列
                for val in wait_lists:
                    await QueueCache.queue_push(QueueCache.VIDEO_JOB, val)
            return None
        except Exception as e:
            logger.error(f"crontab _distribute_video: " + str(e))

    @classmethod
    async def _distribute_music(cls):
        try:
            # 队列中元素数量
            queue_data = await QueueCache.queue_peek(QueueCache.MUSIC_JOB)
            if len(queue_data) > 10:
                return None

            # 查询超时未完成 (10分钟未处理完)
            expire_time: int = int(10 * 60)
            current_time: int = int(time.time())
            err_records = await (MusicRecordsModel
                                 .filter(status=MusicEnum.STATUS_ING)
                                 .filter(start_time__lte=(current_time - expire_time))
                                 .filter(is_delete=0)
                                 .limit(100)
                                 .order_by("id")
                                 .values("id", "user_id", "use_points", "create_time"))

            # 处理超时未完成
            for record in err_records:
                # 已超时
                await MusicRecordsModel.filter(id=record["id"]).update(
                    fail_reason="任务响应超时",
                    status=MusicEnum.STATUS_FAIL,
                    update_time=int(time.time())
                )
                # 回退积分
                if record["use_points"] > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record["user_id"],
                        change_type=WalletEnum.POINTS_INC_MUSIC_FAIL,
                        change_amount=record["points"],
                        project="AI音乐"
                    )

            # 获取等待中数据
            wait_records = await (MusicRecordsModel
                                  .filter(is_delete=0)
                                  .filter(status__in=[MusicEnum.STATUS_WAIT, MusicEnum.STATUS_ING])
                                  .limit(50)
                                  .order_by("id")
                                  .values(str("id")))

            # 处理待视频任务
            wait_lists = [item["id"] for item in wait_records]
            if wait_lists:
                # 更新数据排队中
                await (MusicRecordsModel
                       .filter(is_delete=int(0))
                       .filter(id__in=wait_lists)
                       .update(
                            status=MusicEnum.STATUS_ING,
                            start_time=int(time.time()),
                            update_time=int(time.time())
                       ))
                # 将数据推入队列
                for val in wait_lists:
                    await QueueCache.queue_push(QueueCache.MUSIC_JOB, val)
            return None
        except Exception as e:
            logger.error(f"crontab _distribute_music: " + str(e))
